#-------------------------------------------------------------------------------
# Name:        taskfile
# Purpose:
#
# Author:      SUPER_紫电  QQ:911344755
#
# Created:     07/09/2022
# Copyright:   (c) SUPER_紫电 2022 All rights reserved.
# Licence:     <Modified BSD License>
#-------------------------------------------------------------------------------

import os, sys, struct
from time import time, sleep
from ctypes import addressof, byref, cast, c_buffer, memmove, memset, sizeof
from ata import *
from memorymap import *
import ioeng, pci

BIT0 = AMNF = ERR  = 0x01
BIT1 = TONF = INDX = 0x02
BIT2 = ABRT = CORR = 0x04
BIT3 = MCR  = DRQ  = 0x08
BIT4 = IDNF = DSC  = 0x10
BIT5 = MC   = DF   = 0x20
BIT6 = UNC  = DRDY = 0x40
BIT7 = BBK  = BSY  = 0x80

ata_controller = []
ata_port = {'base': 0x0, 'alt': 0x0, 'dev': 0x40,}
ata_dev = {'master': 0x40, 'slave': 0x50,}

ata_reg = {
    'data': {'address': 0x0, 'value': 0x0}, # Data Register
    'fe'  : {'address': 0x0, 'value': 0x0}, # 1 - Features | Error Register
    'sc'  : {'address': 0x0, 'value': 0x0}, # 2 - Sector Count | Interrupt Reason Register
    'low' : {'address': 0x0, 'value': 0x0}, # 3 - LBA Low Register | Sector Number Register
    'mid' : {'address': 0x0, 'value': 0x0}, # 4 - LBA Mid/Byte Count Low Register | Cylinder Low Register
    'high': {'address': 0x0, 'value': 0x0}, # 5 - LBA High/Byte Count High Register | Cylinder High Register
    'dev' : {'address': 0x0, 'value': 0x0}, # 6 - Device Register
    'cmd' : {'address': 0x0, 'value': 0x0}, # 7 - Command | Status Register
    'alt' : {'address': 0x0, 'value': 0x0}, # Device Control | Alternate Status Register
}

drv_initialized = False
cur_error  = 0xEE
cur_status = 0xEE
ident_dev_data = IDENTIFY_DEVICE_DATA()

def inp(port): return ioeng.inp(port)
def inpw(port): return ioeng.inpw(port)
def inpd(port): return ioeng.inpd(port)
def outp(port, value): return ioeng.outp(port, value)
def outpw(port, value): return ioeng.outpw(port, value)
def outpd(port, value): return ioeng.outpd(port, value)

def init_drv():
    global drv_initialized
    if not create_shared_memory():
        return False
    pid = os.getpid()
    drv_initialized = ioeng.init_drv(pid)
    return drv_initialized

def deinit_drv():
    close_shared_memory()
    if drv_initialized:
        ioeng.deinit_drv()

def enum_ata_controller(ata_controller):
    bar = pci.PCICFG_REG_BAR0
    for bus in range(256):
        for dev in range(32):
            for func in range(8):
                address = pci.pci_bus_dev_func(bus, dev, func)
                vid_value = ioeng.read_pci_config_dword(address, 0x00)
                if (vid_value == 0xFFFFFFFF) or (vid_value == 0):
                    continue
                class_value = ioeng.read_pci_config_dword(address, 0x08)
                header_type = ioeng.read_pci_config_byte(address, 0x0E)
                if class_value >> 24 & 0xFF != 1:
                    continue
                if header_type != 0:
                    continue
                for i in range(3):
                    base = ioeng.read_pci_config_dword(address, bar + i * 8)
                    alt = ioeng.read_pci_config_dword(address, bar + i * 8 + 4)
                    if(base&1) and (alt&1) and (base>0x170) and (alt<0xFFFF):
                        ata_port = {}
                        ata_port['base'] = base - 1
                        ata_port['alt' ] = alt + 1
                        ata_port['dev' ] = 0x40
                        ata_controller.append(ata_port)
                        ata_port = {}
                        ata_port['base'] = base - 1
                        ata_port['alt' ] = alt + 1
                        ata_port['dev' ] = 0x50
                        ata_controller.append(ata_port)

def set_ata_controller():
    global ata_controller
    ata_controller.clear()
    enum_ata_controller(ata_controller)

def device_select(device):
    global ata_port
    global ata_reg
    ata_reg['dev']['value'] = device
    ata_port['dev'] = device

def set_reg_addr(base, alt, dev=0x40):
    global ata_port
    global ata_reg
    ata_port['base'] = base
    ata_port['alt' ] = alt
    device_select(dev)
    ata_reg['data']['address'] = base
    ata_reg['fe'  ]['address'] = base + 1
    ata_reg['sc'  ]['address'] = base + 2
    ata_reg['low' ]['address'] = base + 3
    ata_reg['mid' ]['address'] = base + 4
    ata_reg['high']['address'] = base + 5
    ata_reg['dev' ]['address'] = base + 6
    ata_reg['cmd' ]['address'] = base + 7
    ata_reg['alt' ]['address'] = alt
    ioeng.set_reg_addr(base, alt, dev)

def get_reg_value():
    global ata_reg
    global cur_error
    global cur_status
    ata_reg['fe'  ]['value'] = inp(ata_reg['fe'  ]['address'])
    ata_reg['sc'  ]['value'] = inp(ata_reg['sc'  ]['address'])
    ata_reg['low' ]['value'] = inp(ata_reg['low' ]['address'])
    ata_reg['mid' ]['value'] = inp(ata_reg['mid' ]['address'])
    ata_reg['high']['value'] = inp(ata_reg['high']['address'])
    ata_reg['dev' ]['value'] = inp(ata_reg['dev' ]['address'])
    ata_reg['cmd' ]['value'] = inp(ata_reg['cmd' ]['address'])
    ata_reg['alt' ]['value'] = inp(ata_reg['alt' ]['address'])
    cur_error  = ata_reg['fe' ]['value']
    cur_status = ata_reg['alt']['value']

def reg_error(bit):
    return inp(ata_reg['fe']['address']) & bit

def reg_status(bit):
    return inp(ata_reg['alt']['address']) & bit

def eob():
    return reg_status(ERR) or reg_status(BSY)

def bit(bit, value):
    buffer = bytearray(64)
    data = value
    for i in range(64):
        buffer[i] = data & 1
        data >>= 1
    return buffer[bit] == 1

def rstat_bit(bit=0):
    #reg_bit = [ERR, INDX, CORR, DRQ, DSC, DF, DRDY, BSY]
    reg_bit = [1, 2, 4, 8, 0x10, 0x20, 0x40, 0x80]
    return inp(ata_reg['alt']['address']) & reg_bit[bit]

def rstat_bitclr(bit=0):
    #reg_bit = [1, 2, 4, 8, 0x10, 0x20, 0x40, 0x80]
    #return (inp(ata_reg['alt']['address']) & reg_bit[bit]) != reg_bit[bit]
    return rstat_bit(bit) == False

def check_drq(timeout=10):
    if reg_status(DRQ):
        return True
    result = False
    start = time()
    while True:
        if reg_status(DRQ):
            result = True
            break
        if time() - start > timeout:
            break
        sleep(0.001)
    return result

def wait_not_busy(timeout=10):
    result = True
    stat = inp(ata_reg['alt']['address'])
    if (stat & BSY):
        start = time()
        while (stat & BSY):
            if time() - start > timeout:
                result = False
                break
            sleep(0.001)
            stat = inp(ata_reg['alt']['address'])
    return result

def eto(timeout):
    return reg_status(ERR) or not wait_not_busy(timeout)

def write_lba(ulba, sc, cmd, fe=0):
    lba = bytearray(6)
    for i in range(6):
        lba[i] = ulba >> 8
        ulba >>= 8
    outp(ata_reg['fe'  ]['address'], fe >> 8)
    outp(ata_reg['sc'  ]['address'], sc >> 8)
    outp(ata_reg['low' ]['address'], lba[3])
    outp(ata_reg['mid' ]['address'], lba[4])
    outp(ata_reg['high']['address'], lba[5])
    outp(ata_reg['fe'  ]['address'], fe & 0xFF)
    outp(ata_reg['sc'  ]['address'], sc & 0xFF)
    outp(ata_reg['low' ]['address'], lba[0])
    outp(ata_reg['mid' ]['address'], lba[1])
    outp(ata_reg['high']['address'], lba[2])
    outp(ata_reg['dev' ]['address'], ata_reg['dev']['value'])
    outp(ata_reg['cmd' ]['address'], cmd)

def lba_to_bytes(ulba):
    lba = bytearray(6)
    for i in range(5, -1, -1):
        lba[i] = ulba & 0xFF
        ulba >>= 8
    return lba

def get_lba():
    lba = bytearray(6)
    lba[3] = inp(ata_reg['high']['address'])
    lba[4] = inp(ata_reg['mid' ]['address'])
    lba[5] = inp(ata_reg['low' ]['address'])
    outp(ata_reg['alt']['address'], 0x80)
    lba[0] = inp(ata_reg['high']['address'])
    lba[1] = inp(ata_reg['mid' ]['address'])
    lba[2] = inp(ata_reg['low' ]['address'])
    ulba = 0
    for i in range(5):
        ulba += lba[i]
        ulba <<= 8
    return ulba + lba[5]

def get_max_lba():
    if not ata_cmd(cmd=0x27):
        return 0
    return get_lba()

def set_max_lba(ulba):
    if ulba > get_max_lba():
        return False
    lba = lba_to_bytes(ulba)
    outp(ata_reg['low' ]['address'], lba[2])
    outp(ata_reg['low' ]['address'], lba[5])
    outp(ata_reg['mid' ]['address'], lba[1])
    outp(ata_reg['mid' ]['address'], lba[4])
    outp(ata_reg['high']['address'], lba[0])
    outp(ata_reg['high']['address'], lba[3])
    outp(ata_reg['dev' ]['address'], ata_reg['dev']['value'])
    outp(ata_reg['cmd' ]['address'], 0x37)
    return ulba == get_lba()

def read_verify_lba(ulba, sc, timeout=10):
    write_lba(ulba, sc, READ_VERIFY_SECTORS_EXT)
    start = time()
    while True:
        stat = inp(ata_reg['alt']['address'])
        while (stat & BSY):
            if time() - start > timeout:
                cmd_recal()
                return -6
            stat = inp(ata_reg['alt']['address'])
        if (stat & INDX):
            cmd_recal()
            ulba -= 1
            return -15
        if stat == 0x10:
            return -16
        if stat == 0x50 or stat == 0x51:
            break
    erro = inp(ata_reg['fe']['address'])
    if (erro & AMNF):
        cmd_recal()
        return -7
    if (stat & ERR):
        result = -8  if (erro & TONF) else -1
        result = -9  if (erro & ABRT) else -1
        result = -10 if (erro & IDNF) else -1
        result = -11 if (erro & UNC ) else -1
        result = -12 if (erro & BBK ) else -1
        ulba = get_lba()
        cmd_recal()
    else:
        result = time() - start
        ulba = ulba + sc - 1
    return result

def write_uncorrectable(ulba, sc):
    write_lba(ulba, sc, WRITE_UNCORRECTABLE_EXT, 0xA5)

def send_ata_cmd(fe, sc, low, mid, high, dev, cmd, timeout=10):
    outp(ata_reg['fe'  ]['address'], fe   >> 8)
    outp(ata_reg['sc'  ]['address'], sc   >> 8)
    outp(ata_reg['low' ]['address'], low  >> 8)
    outp(ata_reg['mid' ]['address'], mid  >> 8)
    outp(ata_reg['high']['address'], high >> 8)
    outp(ata_reg['fe'  ]['address'], fe   & 0xFF)
    outp(ata_reg['sc'  ]['address'], sc   & 0xFF)
    outp(ata_reg['low' ]['address'], low  & 0xFF)
    outp(ata_reg['mid' ]['address'], mid  & 0xFF)
    outp(ata_reg['high']['address'], high & 0xFF)
    outp(ata_reg['dev' ]['address'], dev)
    outp(ata_reg['cmd' ]['address'], cmd)
    if not wait_not_busy(timeout):
        return False
    return rstat_bitclr(0)

def ata_cmd(fe=0, sc=0, low=0, mid=0, high=0, cmd=0, timeout=10):
    dev = ata_reg['dev']['value']
    return send_ata_cmd(fe, sc, low, mid, high, dev, cmd, timeout)

def soft_reset():
    outp(ata_reg['alt']['address'], 4)
    sleep(0.001)
    outp(ata_reg['alt']['address'], 0)

def cmd_identify():
    return ata_cmd(cmd=0xEC) and reg_status(DRQ)

def cmd_sleep_drive():
    return ata_cmd(cmd=0xE6)

def cmd_standby():
    return ata_cmd(cmd=0xE0, timeout=3)

def cmd_idle():
    return ata_cmd(cmd=0xE1, timeout=30)

def cmd_wake_drive():
    soft_reset()
    dev= ata_reg['dev']['value']
    outp(ata_reg['dev']['address'], dev)
    outp(ata_reg['cmd']['address'], 0xE0)
    if not wait_not_busy(3):
        return False
    outp(ata_reg['dev']['address'], dev)
    outp(ata_reg['cmd']['address'], 0xE1)
    return wait_not_busy(30)

def cmd_recal():
    return ata_cmd(cmd=0x10, timeout=30)

def cmd_rdy():
    return send_ata_cmd(0,0,0,0,0,0x10,0)

def cmd_set_pwd():
    return ata_cmd(cmd=0xF1) and reg_status(DRQ)

def cmd_unlock_unit():
    return ata_cmd(cmd=0xF2) and reg_status(DRQ)

def cmd_erase_prep():
    return ata_cmd(cmd=0xF3)

def cmd_erase_unit():
    return ata_cmd(cmd=0xF4) and reg_status(DRQ)

def cmd_freeze_lock():
    return ata_cmd(cmd=0xF5)

def cmd_disable_pwd():
    return ata_cmd(cmd=0xF6) and reg_status(DRQ)

def cmd_smart_on():
    return ata_cmd(0xD8,0,0,0x4F,0xC2,0xB0) and reg_status(DRDY)

def read_file(filename, buffer):
    try:
        with open(filename, 'rb') as file:
            while True:
                data = file.read(4)
                if not data:
                    break
                buffer.extend(struct.unpack('@4B', data))
        return True
    except:
        return False

def save_file(filename, buffer):
    try:
        with open(filename, 'wb') as file:
            file.write(buffer)
        return True
    except:
        return False

def read_data(buffer, timeout=10):
    count = len(buffer) // 512
    if count == 0:
        return False
    if not check_drq():
        return False
    if not ioeng.read_data(0, count, timeout):
        return False
    if not wait_not_busy(10):
        return False
    return read_shared_memory(buffer, 0, count*512)

def write_data(buffer, timeout=10):
    count = len(buffer) // 512
    if count == 0:
        return False
    if not write_shared_memory(buffer, 0, count*512):
        return False
    if not check_drq():
        return False
    return ioeng.write_data(0, count, timeout)

def get_identify():
    ident_buffer = bytearray(512)
    get_identify_data(ident_buffer)
    return dec_identify(ident_buffer)

def get_identify_data(buffer):
    global ident_dev_data
    if not cmd_identify():
        return False
    if not read_data(buffer):
        return False
    memset(addressof(ident_dev_data), 0, sizeof(ident_dev_data))
    memmove(addressof(ident_dev_data), byref(c_buffer(bytes(buffer))), len(buffer))
    return True

def get_char(buffer, offset, size):
    temp = ''
    for i in range(0, size, 2):
        temp += chr(buffer[offset + i + 1])
        temp += chr(buffer[offset + i])
    return temp.strip()

def dec_identify(buffer):
    ident = '\nDrive Identify Data'
    ident += '\n------------------------------------------------------------'
    model  = get_char(buffer, 54, 40).replace('\x00', '')
    fwrev  = get_char(buffer, 46, 8 ).replace('\x00', '')
    serial = get_char(buffer, 20, 20).replace('\x00', '')
    ident += '\nModel Number   : ' + model
    ident += '\nFirmware Rev.  : ' + fwrev
    ident += '\nSerial Number  : ' + serial
    lba2 = struct.unpack_from('@LL', buffer, 200)
    lba = lba2[0] + (lba2[1] << 32)
    #ident += '\nUser Sectors   : '+str(lba)+' ('+str(int(lba*512/1000000000))+'G)'
    ident += '\nUser Sectors   : ' + str(lba)
    if lba2[1] == 0:
        ident += ' (' + str(int(lba*512/1000000000)) + 'G)'
    else:
        ident += ' (' + str(int(lba*512/1000000)) + 'T)'
    uflag = buffer[82] + (buffer[83] << 8)
    if uflag & 0x400:
        ident += ' - HPA'
    ident += '\nBuffers        : ' + str(int((buffer[42]+(buffer[43]<<8))/2))+'KB'
    # Security Mode Feature Set
    uflag = buffer[164] + (buffer[165] << 8)
    if uflag & 2:
        ident += '\nFeature Sets   : (0x%04X)security supported' % uflag
    else:
        ident += '\nFeature Sets   : (0x%04X)security unsupported' % uflag
    # Security Status
    uflag = buffer[256] + (buffer[257] << 8)
    if uflag & 1:
        ident += '\nSecurity Status: (0x%04X)security supported' % uflag
        if uflag & 2:
            ident += ', enabled'
            if uflag & 256:
                ident += ', level is maximum'
            else:
                ident += ', level is high'
        elif uflag & 4:
            ident += ', locked'
        elif uflag & 8:
            ident += ', frozen'
        elif uflag & 0x10:
            ident += ', count has expired'
        # Time required for Security Erase Unit Completion
        if uflag & 0x20:
            ident += '\nEnhanced Erase : '
            value = buffer[180] + (buffer[181] << 8)
            if value == 0:
                ident += 'value not specified'
            elif value == 255:
                ident += '>508 minutes'
            else:
                ident += str(value * 2) + ' minutes'
        ident += '\nNormal Erase   : '
        value = buffer[178] + (buffer[179] << 8)
        if value == 0:
            ident += 'value not specified'
        elif value == 255:
            ident += '>508 minutes'
        else:
            ident += str(value * 2) + ' minutes'
    # Display Automatic Acoustic Managment
    ident += '\nCurrent AAM Lvl: '
    uflag = buffer[166] + (buffer[167] << 8)
    if uflag & 0x200:
        uflag = buffer[188]
        ident += str(uflag) + ' '
        if uflag == 0xFF:
            ident += '(Reserved)'
        elif uflag == 0xFE:
            ident += '(Performance)'
        elif uflag >= 0x81:
            ident += '(Intermediate)'
        elif uflag == 0x80:
            ident += '(Quiet Mode)'
        elif uflag >= 1:
            ident += '(Retired)'
        else:
            ident += '(Reserved)'
        uflag = buffer[172] + (buffer[173] << 8)
        if uflag & 0x200 == 0:
            ident += ', AAM is disabled'
        else:
            ident += ', AAM is enabled'
    else:
        ident += 'AAM not supported'
    # Display UDMA Mode
    ident += '\nUdma Mode      : '
    uflag = buffer[176]
    if uflag == 0xFF:
        ident += 'Disabled'
    else:
        for i in range(9):
            if uflag & 1:
                ident += str(i) + ' '
            uflag >>= 1
    # Display Cache State
    ident += '\nCache State    : '
    uflag = buffer[170]
    if uflag & 64:
        ident += 'RCache On, '
    else:
        ident += 'RCache Off, '
    if uflag & 32:
        ident += 'WCache On'
    else:
        ident += 'WCache Off'
    # Display Smart State
    ident += '\nSmart Enabled  : '
    if uflag & 1:
        ident += 'On'
    else:
        ident += 'Off'
    return ident
